from torch import FloatTensor
from torch import squeeze
from torch_geometric.nn import MessagePassing
import torch.nn as nn
import torch.nn.functional as F


class MLPLayer(MessagePassing):
    def __init__(self,  in_features, out_features, device="cpu"):
        """MPLayer  普通的MessagePassinglayer

        Args:
            in_feature ([type]): 传入的features
            out_features ([type]): 传出的features
            device ([type]): device
        """
        super(MLPLayer, self).__init__(aggr='add')  # "Add" aggregation (Step 5).
        self.linear = nn.Linear(in_features, out_features)
        self.device = device

        
    
    def forward(self, x, edge_index):
        x = self.linear(x)
        return self.propagate(edge_index, x=x)


if __name__ == '__main__':
    import torch
    import numpy
    from torch_geometric.data import Data, InMemoryDataset, Dataset
    from torch_geometric.loader import DataLoader
    class TestDataset(Dataset):
        def __init__(self):
            self.node=torch.rand(10,10,16)
            self.edge=torch.randint(0,16,[10,2,10])
            
        def __getitem__(self, idx):
            return self.node[idx],self.edge[idx]
        
        def __len__(self) -> int:
            return 10
        
    dataloader=DataLoader(TestDataset(),2,True)
    model = MLPLayer( in_features=16, out_features=25, device="cpu")
    for b in dataloader:
        node,edge=b
        
        model(node,edge)
        
        #print(node.size())
        #print(edge.size())
    
    # b = model(a,torch.randint(0,16,[3,2,10]) )
    # print(b.size())
    # print("")